# @file    : test2
# @time    : 2025/4/17
# @author  : yongpeng.yao
# @desc    :
import os

import pandas as pd
from dotenv import load_dotenv
from langchain.schema.output_parser import StrOutputParser
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import DataFrameLoader
from langchain_community.vectorstores import Chroma
from langchain_core.prompts import ChatPromptTemplate
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_openai import ChatOpenAI

load_dotenv()


# 1. 读取Excel文件
def read_excel_to_dataframe(file_path, sheet_name=0):
    """
    读取Excel文件到Pandas DataFrame
    :param file_path: Excel文件路径
    :param sheet_name: 工作表名或索引，默认为第一个工作表
    :return: Pandas DataFrame
    """
    try:
        df = pd.read_excel(file_path, sheet_name=sheet_name, engine='openpyxl')
        print(f"成功读取Excel文件，共{len(df)}行数据")
        return df
    except Exception as e:
        print(f"读取Excel文件失败: {e}")
        return None


# 2. 预处理数据并创建向量数据库
def create_vector_db(df, text_columns=None, persist_directory="chroma_db"):
    """
    从DataFrame创建向量数据库
    :param df: 包含文本数据的DataFrame
    :param text_columns: 包含文本数据的列名列表
    :param persist_directory: 向量数据库存储目录
    :return: 向量数据库对象
    """
    if not text_columns:
        text_columns = df.columns
    # 合并文本列
    df['combined_text'] = df[text_columns].apply(lambda x: ' '.join(x.dropna().astype(str)), axis=1)

    # 使用DataFrameLoader加载数据
    loader = DataFrameLoader(df, page_content_column='combined_text')
    documents = loader.load()

    # 文本分割
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=1000,
        chunk_overlap=200
    )
    texts = text_splitter.split_documents(documents)

    # 创建嵌入模型
    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")

    # 创建向量数据库
    vectordb = Chroma.from_documents(
        documents=texts,
        embedding=embeddings,
        persist_directory=persist_directory
    )
    print("向量数据库创建完成并持久化存储")
    return vectordb


# 3. 使用DeepSeek模型分析总结
def analyze_with_deepseek(query, vectordb, top_k=3):
    """
    使用DeepSeek模型分析总结
    :param query: 查询文本
    :param vectordb: 向量数据库
    :param top_k: 返回最相关的文档数量
    :return: 分析总结结果
    """
    # 从向量数据库检索相关文档
    relevant_docs = vectordb.similarity_search(query, k=top_k)
    context = "\n\n".join([doc.page_content for doc in relevant_docs])

    # 初始化DeepSeek API (假设存在)
    deepseek = ChatOpenAI(model='deepseek-r1',
                          api_key=os.getenv("DASHSCOPE_API_KEY"),
                          base_url="https://dashscope.aliyuncs.com/compatible-mode/v1"
                          )

    # 构造提示词
    prompt = """请根据以下上下文信息回答问题或进行总结分析：

    上下文：
    {context}

    问题或任务：{query}

    请提供详细的分析和总结，重点突出关键信息和见解。"""
    chat_prompt = ChatPromptTemplate.from_template(prompt)
    # 调用DeepSeek模型
    output_parser = StrOutputParser()
    chain = chat_prompt | deepseek | output_parser
    return chain.invoke({"context": context, "query": query})


# 主函数
def main():
    # 配置参数
    excel_path = "./data/translated2.xlsx"  # 替换为你的Excel文件路径
    sheet_name = 0  # 可以使用工作表名或索引
    text_columns = None  # 替换为你的Excel中包含文本的列名
    persist_directory = "excel_chroma_db"

    # 1. 读取Excel数据
    df = read_excel_to_dataframe(excel_path, sheet_name)
    if df is None:
        return

    # 2. 创建向量数据库
    vectordb = create_vector_db(df, text_columns, persist_directory)

    # 3. 交互式查询和分析
    while True:
        print("\n请输入你的查询或分析需求（输入'exit'退出）：")
        user_query = input().strip()

        if user_query.lower() == 'exit':
            break

        if not user_query:
            print("查询不能为空，请重新输入")
            continue

        # 使用DeepSeek分析
        print("\n分析结果：")
        result = analyze_with_deepseek(user_query, vectordb)
        print(result)


if __name__ == "__main__":
    main()
